{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hackathon 9\n",
    "\n",
    "Written by Eleanor Quint\n",
    "with images sourced from [Chris Olah's blog post on LSTM](http://colah.github.io/posts/2015-08-Understanding-LSTMs/).\n",
    "\n",
    "Topics:\n",
    "- TensorFlow RNNs and Cells\n",
    "- LSTM\n",
    "\n",
    "In today's demo, we'll teach an RNN how to speak English.\n",
    "\n",
    "This is all setup in a IPython notebook so you can run any code you want to experiment with. Feel free to edit any cell, or add some to run your own code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# We'll start with our library imports...\n",
    "from __future__ import print_function\n",
    "\n",
    "import random\n",
    "import os  # to work with file paths\n",
    "\n",
    "import tensorflow as tf         # to specify and run computation graphs\n",
    "from tensorflow import keras\n",
    "import numpy as np              # for numerical operations taking place outside of the TF graph\n",
    "import matplotlib.pyplot as plt # to draw plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### RNN/LSTM theory recap\n",
    "\n",
    "Recurrent neural networks (RNNs) are computation graphs with loops (i.e., not directed acyclic graphs). Because the backpropagation algorithm only works with DAGs, we have to unroll the RNN through time. Tensorflow provides code that handles this automatically.\n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/RNN-unrolled.png\" width=\"80%\">\n",
    "\n",
    "\n",
    "The most common RNN unit is the LSTM, depicted below:\n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-chain.png\" width=\"80%\">\n",
    "\n",
    "We can see that each unit takes 3 inputs and produces 3 outputs, two which are forwarded to the same unit at the next timestep and one true output, $h_t$ depicted coming out of the top of the cell.\n",
    "\n",
    "The upper right output going to the next timestep is the cell state. It carries long-term information between cells, and is calculated as: \n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-focus-C.png\" width=\"80%\">\n",
    "\n",
    "where the first term uses the forget gate $f_t$ to decide to scale the previous state (potentially making it smaller to \"forget\" it), and the second term is the product of the update gate $i_t$ and the state update $\\tilde{C}_t$. Each of the forget and update gates are activated with sigmoid, so their range is (0,1).\n",
    "\n",
    "The true output and the second, lower output on the diagram are calculated by the output gate:\n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-focus-o.png\" width=\"80%\">\n",
    "\n",
    "First, $o_t$ is calculated from the output of the previous timestep concatenated with the current input, but then it's mixed with the cell state to get the true output. Passing on this output to the next timestep as the hidden state gives the unit a kind of short term memory.\n",
    "\n",
    "(Images sourced from [Colah's Blog](http://colah.github.io/posts/2015-08-Understanding-LSTMs/))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Today, we're going to teach a recurrent model how to speak English by starting from a sequence of words and asking the model to predict what the next word should be. And what better way to learn English than by learning to talk like an angry media reviewer on the internet? We'll be using the IMDB review corpus for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "imdb = keras.datasets.imdb\n",
    "\n",
    "(train_data, train_labels), (test_data, test_labels) = imdb.load_data(num_words=10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data is encoded as a list of integers. We'll get the mapping of words to integer values in order to be able to translate the words back and forth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# A dictionary mapping words to an integer index\n",
    "word_index = imdb.get_word_index()\n",
    "\n",
    "# The first indices are reserved\n",
    "word_index = {k:(v+3) for k,v in word_index.items()} \n",
    "word_index[\"<PAD>\"] = 0\n",
    "word_index[\"<START>\"] = 1\n",
    "word_index[\"<UNK>\"] = 2  # unknown\n",
    "word_index[\"<UNUSED>\"] = 3\n",
    "\n",
    "reverse_word_index = dict([(value, key) for (key, value) in word_index.items()])\n",
    "\n",
    "def decode_review(text):\n",
    "    return ' '.join([reverse_word_index.get(i, '?') for i in text if i != 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we'll pad the text sequences so that they're all of identical length and easier to work with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEQ_LEN = 256\n",
    "\n",
    "train_data = keras.preprocessing.sequence.pad_sequences(train_data,\n",
    "                                                        value=word_index[\"<PAD>\"],\n",
    "                                                        padding='post',\n",
    "                                                        maxlen=SEQ_LEN)\n",
    "train_targets = keras.preprocessing.sequence.pad_sequences(train_data[:,1:],\n",
    "                                                        value=word_index[\"<START>\"],\n",
    "                                                        padding='post',\n",
    "                                                        maxlen=SEQ_LEN)\n",
    "\n",
    "test_data = keras.preprocessing.sequence.pad_sequences(test_data,\n",
    "                                                       value=word_index[\"<PAD>\"],\n",
    "                                                       padding='post',\n",
    "                                                       maxlen=SEQ_LEN)\n",
    "test_targets = keras.preprocessing.sequence.pad_sequences(test_data[:,1:],\n",
    "                                                        value=word_index[\"<START>\"],\n",
    "                                                        padding='post',\n",
    "                                                        maxlen=SEQ_LEN)\n",
    "\n",
    "print(\"Training data has shape {}\".format(train_data.shape))\n",
    "print(\"Training targets has shape {}\".format(train_targets.shape))\n",
    "print(\"Testing data has shape {}\".format(train_data.shape))\n",
    "print(\"Testing targets has shape {}\".format(test_targets.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize some of the data\n",
    "idx = random.randrange(train_data.shape[0])\n",
    "print(train_data[idx])\n",
    "print(decode_review(train_data[idx]))\n",
    "print(train_targets[idx])\n",
    "print(decode_review(train_targets[idx]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each datum is a string of up to 256 successive words from the corpus, and the target is a similar window, but shifted forward by one word. This is setup to train the model to, given a few preceding words, predict what the next word in the sequence will be.\n",
    "\n",
    "Initially, in the data each word in the sequence is represented as an integer (notice the shape). This discrete representation fails to capture any semantic relationships between words. I.e., the model wouldn't know that \"crimson\" and \"scarlet\" are more similar than \"red\" and \"blue\". The solution is to learn an word embedding as the first part of the model to transform each integer into a relatively small, dense vector (as compared to a one-hot). Then, similar words will train to have similar embeddings.\n",
    "\n",
    "We'll use [tf.nn.embedding_lookup](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup) to do this which we provide a (usually trainable) VOCAB_SIZE x EMBEDDING_SIZE matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = max(word_index.values())\n",
    "EMBEDDING_SIZE = int(np.sqrt(VOCAB_SIZE))\n",
    "print(\"Vocab size is {} and is embedded into {} dimensions\".format(VOCAB_SIZE, EMBEDDING_SIZE))\n",
    "\n",
    "# setup input and embedding\n",
    "input_ph = tf.placeholder(tf.int32, [None, 256])\n",
    "target_ph = tf.placeholder(tf.int32, [None, 256])\n",
    "\n",
    "with tf.variable_scope(\"embedding\", reuse=tf.AUTO_REUSE):\n",
    "    embedding_matrix = tf.get_variable('embedding_matrix', dtype=tf.float32, shape=[VOCAB_SIZE, EMBEDDING_SIZE],\n",
    "                                       trainable=True)\n",
    "word_embeddings = tf.nn.embedding_lookup(embedding_matrix, input_ph)\n",
    "print(\"The output of the word embedding: \" + str(word_embeddings))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to declare an architecture that looks like this (replacing GRU with LSTM and char embedding with word embedding).\n",
    "\n",
    "<img src=\"https://tensorflow.org/tutorials/sequences/images/text_generation_training.png\" width=\"80%\">\n",
    "\n",
    "TensorFlow separates the declaration of [RNNCells](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/RNNCell) from the [RNNs](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn) that run them. In the code below, we declare an [LSTM cell](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/BasicLSTMCell), and create tensors for the inputs to the first unit. We use zeros for the initial hidden state and current state, but it's also possible to declare trainable variables for these as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSTM_SIZE = 200 # number of units in the LSTM layer, this number taken from a \"small\" language model\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "lstm_cell = tf.contrib.rnn.LSTMCell(LSTM_SIZE)\n",
    "\n",
    "# Initial state of the LSTM memory.\n",
    "initial_state = lstm_cell.zero_state(BATCH_SIZE, tf.float32)\n",
    "print(\"Initial state of the LSTM: \" + str(initial_state))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we'll pass the newly declared cell and the training sequence of word embeddings to [tf.nn.dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn) as the inputs over time to the LSTM. `dynamic_rnn` runs an `RNNCell` using an internal `while` loop, and returns the sequence of outputs from the LSTM at each timestep and the final state of the LSTM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup RNN\n",
    "outputs, state = tf.nn.dynamic_rnn(lstm_cell, word_embeddings,\n",
    "                                   initial_state=initial_state,\n",
    "                                   dtype=tf.float32)\n",
    "print(\"The outputs over all timesteps: \"+ str(outputs))\n",
    "print(\"The final state of the LSTM layer: \" + str(state))\n",
    "logits = tf.layers.dense(outputs, VOCAB_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And to calculate the loss between two sequences, we'll import a function from [tf.contrib.seq2seq](https://www.tensorflow.org/api_docs/python/tf/contrib/seq2seq) called [sequence_loss](https://www.tensorflow.org/api_docs/python/tf/contrib/seq2seq/sequence_loss). It calculates the weighted cross-entropy loss between the first two arguments, and the third argument provides weights for averaging. We weight uniformly here, but weights could also be calculated based on where in the sequence the target is (e.g., penalize less earlier in the sequence, but more later) or based on the content of the target (e.g., low weight on guessing articles correctly and larger weight on getting nouns and verbs correct).\n",
    "\n",
    "We'll optimize using TensorFlow's [RMSProp](https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer) optimizer, which requires an explicit learning rate, but otherwise as usual. We switch from the Adam optmizer because we don't want the adaptive learning rate feature, which can interact badly with the recurrent gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "loss = tf.contrib.seq2seq.sequence_loss(\n",
    "    logits,\n",
    "    target_ph,\n",
    "    tf.ones([BATCH_SIZE, SEQ_LEN], dtype=tf.float32), # we'll use uniform weight over timesteps\n",
    "    average_across_timesteps=True,\n",
    "    average_across_batch=True)\n",
    "\n",
    "optimizer = tf.train.RMSPropOptimizer(LEARNING_RATE)\n",
    "train_op = optimizer.minimize(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we'll create a Session, initialize the variables, and run the train op once. This model is relatively heavyweight, so we don't want to optimize it on the login node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = tf.Session()\n",
    "session.run(tf.global_variables_initializer())\n",
    "\n",
    "# we'll just run one step, omitting the usual epoch code\n",
    "_ = session.run(train_op, feed_dict={input_ph: train_data[idx:idx + BATCH_SIZE], target_ph: train_targets[idx:idx + BATCH_SIZE]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll try generating from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# this value is to limit load on the login node, must be <= SEQ_LEN\n",
    "GEN_LEN = 10\n",
    "\n",
    "# generated_sequence needs to be a valid input\n",
    "generated_sequence = np.array([[word_index[\"<START>\"]] + [word_index[\"<PAD>\"]]*(SEQ_LEN - 1)] * BATCH_SIZE) \n",
    "for idx in range(1, GEN_LEN):\n",
    "    logits_val = session.run(logits, feed_dict={input_ph: generated_sequence})\n",
    "    generated_words = np.array([np.argmax(l_val[idx]) for l_val in logits_val])\n",
    "    generated_sequence[:,idx] = generated_words\n",
    "    print(decode_review(generated_sequence[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hackathon 9 Exercise\n",
    "\n",
    "Create a 2 layer [LSTMCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/LSTMCell) with an [Attention Wrapper](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/AttentionCellWrapper). Your code should use `train_ph` inputs as above and your code should finish with the `loss` tensor. This should be pretty straightforward with the TensorFlow documentation.\n",
    "\n",
    "This model is very large and trains for a long time, so please don't try to optimize it in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TensorFlow 1.12 (py36)",
   "language": "python",
   "name": "tensorflow-1.12-py36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
